# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: chunhui

import glob
import json
import os

from absl import app, flags

FLAGS = flags.FLAGS


def merge(exist_path, name, config):
  with open(exist_path) as f:
    driver = json.load(f)
  driver['feed']['exchanges'][name] = config
  return driver


def temp_replacement(symbol):
  return symbol.replace('BCHABC', 'BCHN')


def gen_config(
    template_driver_path,
    output_dir,
    pi_root=None,
    feed_root=None,
    date=None,
    mea_equals=None,
    mea_contains=None,
    worker=None):
  pi_root = pi_root or FLAGS.pi_root
  feed_root = feed_root or FLAGS.feed_root
  mea_equals = mea_equals if mea_equals is not None else FLAGS.mea_equals
  mea_contains = mea_contains if mea_contains is not None else FLAGS.mea_contains
  worker = worker if worker is not None else FLAGS.worker
  date = date or FLAGS.date
  with open(template_driver_path, 'r') as f:
    print(template_driver_path)
    driver = json.load(f)
  machine = driver['feed']['common']['archive']['machine']
  exchanges = driver['feed']['exchanges']
  for key, config in exchanges.items():
    quote_patt = ".*"
    expiry = ".*"
    if "quotes" in config:
      quote_patt = "|".join(config["quotes"])
      del config["quotes"]
    if "expires" in config:
      expiry = "|".join(config["expires"])
      del config["expires"]
    config["products"] = {"norm_format": f".*({quote_patt}).*({expiry}).*"}
    if "groups" in config:
      config["products"] = {"groups": config['groups']}
      del config['groups']
  exist_meas = {}
  for name, exchange in driver['feed']['exchanges'].items():
    mea = exchange['mea']
    if mea_equals and mea_equals.lower() != mea.lower():
      continue
    if mea_contains and mea_contains.lower() not in mea.lower():
      continue
    output_path = os.path.join(output_dir, f'{mea}.json')
    if mea in exist_meas and os.path.exists(output_path):
      current_driver = merge(output_path, name, exchange)
    else:
      current_driver = {
          "feed": {
              "common": driver['feed']['common'],
              "exchanges": {
                  name: exchange
              }
          }
      }
    exist_meas[mea] = 1
    if worker:
      driver['feed']['common']['archive']['worker'] = worker
    with open(output_path, "w") as wf:
      json.dump(current_driver, wf, indent=2)
  return exist_meas.keys()


def main(_):
  assert FLAGS.input_dir
  assert FLAGS.output_dir
  assert FLAGS.date
  if not os.path.exists(FLAGS.output_dir):
    os.makedirs(FLAGS.output_dir)
  else:
    if FLAGS.clear_existing:
      for f in glob.glob(os.path.join(FLAGS.output_dir, '*')):
        os.remove(f)
  for dir_name, _, file_names in os.walk(FLAGS.input_dir):
    for file_name in file_names:
      input_path = os.path.join(dir_name, file_name)
      gen_config(input_path, FLAGS.output_dir)


if __name__ == '__main__':
  flags.DEFINE_string('feed_root', '/remote/iosg/coin-2/feed_cache2', '')
  flags.DEFINE_string('pi_root', './data/coin2/product_info', '')
  flags.DEFINE_string('input_dir', '', '')
  flags.DEFINE_string('output_dir', '', '')
  flags.DEFINE_string('date', '', '')
  flags.DEFINE_string('mea_contains', '', '')
  flags.DEFINE_string('mea_equals', '', '')
  flags.DEFINE_string('worker', '', '')
  flags.DEFINE_boolean('clear_existing', True, '')
  app.run(main)
